{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Trends with the GOZCARDS Dataset\n",
    "\n",
    "Here we calculate trends using the GOZCARDS dataset by regressing to the VMR monthly zonal means using seasonal terms in our predictors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "from LOTUS_regression.regression import regress_all_bins\n",
    "from LOTUS_regression.predictors.seasonal import add_seasonal_components\n",
    "from LOTUS_regression.predictors import load_data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The GOZCARDS data is in multiple NetCDF4 files.  Load them all in and concatenate on the time dimension.  Also only take values in the latitude range -60 to 60."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "GOZCARDS_FILES = r'/home/runner/work/lotus-regression/lotus-regression/test_data//GOZCARDS/*.nc4'\n",
    "\n",
    "data = xr.decode_cf(xr.open_mfdataset(GOZCARDS_FILES, combine='nested', concat_dim='time', group='Merged', engine='netcdf4'))\n",
    "\n",
    "data = data.loc[dict(lat=slice(-60, 60))]\n",
    "\n",
    "print(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load some standard predictors and add a constant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "predictors = load_data('pred_baseline_pwlt.csv')\n",
    "\n",
    "print(predictors.columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Currently our predictors have no seasonal dependence.  Add in some seasonal terms with different numbers of Fourier components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "predictors = add_seasonal_components(predictors, {'constant': 4, 'linear_pre': 2, 'linear_post': 2, 'qboA': 2, 'qboB': 2})\n",
    "\n",
    "print(predictors[:5])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Perform the regression and convert the coefficients to percent anomaly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "results = regress_all_bins(predictors, data['average'], tolerance=0.1)\n",
    "\n",
    "# Convert to ~ percent\n",
    "results /= data['average'].mean(dim='time')\n",
    "results *= 100\n",
    "\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import LOTUS_regression.plotting.trends as trends\n",
    "trends.pre_post_with_confidence(results, x='lat', y='lev', ylim=(100, 0.5), log_y=True, figsize=(16, 6),\n",
    "                                x_label='Latitude [$^\\circ$]', y_label='Pressure [hPa]', pre_title='Pre 1997',\n",
    "                                post_title='Post 1997')"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Edit Metadata",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}